import os
import warnings

warnings.filterwarnings('ignore')
import tensorflow as tf
from tensorflow.keras.preprocessing.image import ImageDataGenerator


def get_generator():
    """
    ResNet由于GPU memory限制, 设置batch_size为16
    Alex Net同等条件下可改为128
    :return:
    """
    # 定义数据所在文件夹
    base_dir = 'D:/SystemSoft/GitRepos/cnn-alex-net/archive/birds/'
    train_dir = os.path.join(base_dir, 'train')
    validation_dir = os.path.join(base_dir, 'valid')
    test_dir = os.path.join(base_dir, 'test')

    train_datagen = ImageDataGenerator(rescale=1. / 255,
                                       rotation_range=10,
                                       width_shift_range=0.1,
                                       height_shift_range=0.1,
                                       horizontal_flip=True,
                                       fill_mode='nearest')
    validation_datagen = ImageDataGenerator(rescale=1. / 255,)
    test_datagen = ImageDataGenerator(rescale=1. / 255)
    train_generator = train_datagen.flow_from_directory(train_dir,
                                                        target_size=(227, 227),
                                                        batch_size=16,
                                                        class_mode='categorical')

    validation_generator = validation_datagen.flow_from_directory(validation_dir,
                                                                  target_size=(227, 227),
                                                                  batch_size=16,
                                                                  class_mode='categorical')
    test_generator = test_datagen.flow_from_directory(test_dir,
                                                      target_size=(227, 227),
                                                      batch_size=1,
                                                      class_mode='categorical')
    return train_generator, validation_generator, test_generator
